import cv2
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt


def canny_edge_detection(image_path, low_threshold=100, high_threshold=200):
    """
    对输入图像进行Canny边缘检测
    
    Args:
        image_path: 图像路径
        low_threshold: Canny算子的低阈值
        high_threshold: Canny算子的高阈值
    
    Returns:
        edges: 边缘检测结果
        gray: 灰度图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图像: {image_path}")
    
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 高斯模糊
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # Canny边缘检测
    edges = cv2.Canny(blurred, low_threshold, high_threshold)
    
    return edges, gray

def compare_images(image1_path, image2_path):
    """
    比较两幅图像的边缘检测结果
    
    Args:
        image1_path: 第一幅图像路径
        image2_path: 第二幅图像路径
    """
    # 对两幅图像进行边缘检测
    edges1, gray1 = canny_edge_detection(image1_path)
    edges2, gray2 = canny_edge_detection(image2_path)
    
    # 确保两幅图像大小相同
    if edges1.shape != edges2.shape:
        raise ValueError("两幅图像大小不一致")
    
    # 计算边缘检测结果的差异
    diff = cv2.absdiff(edges1, edges2)
    
    # 计算差异度量
    total_pixels = edges1.shape[0] * edges1.shape[1]
    different_pixels = np.count_nonzero(diff)
    similarity = 1 - (different_pixels / total_pixels)
  
    # 显示结果
    plt.figure(figsize=(15, 10))
    
    plt.subplot(231)
    plt.imshow(gray1, cmap='gray')
    plt.title('src1')
    
    plt.subplot(232)
    plt.imshow(gray2, cmap='gray')
    plt.title('src2')
    
    plt.subplot(233)
    plt.imshow(diff, cmap='gray')
    plt.title('diff')
    
    plt.subplot(234)
    plt.imshow(edges1, cmap='gray')
    plt.title('edge1')
    
    plt.subplot(235)
    plt.imshow(edges2, cmap='gray')
    plt.title('edge2')
    
    plt.subplot(236)
    plt.text(0.5, 0.5, f'similarity: {similarity:.2%}', 
             horizontalalignment='center',
             verticalalignment='center',
             fontsize=12)
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()

# if __name__ == "__main__":
#     image1_path = "./image1.jpg"
#     image2_path = "./image3.jpg"
    
# compare_images(image1_path, image2_path)


# 读取两幅图像
ref_image = cv2.imread("./image1.jpg", cv2.IMREAD_GRAYSCALE)
target_image = cv2.imread("./image3.jpg", cv2.IMREAD_GRAYSCALE)

# 高斯模糊
blurred_ref = cv2.GaussianBlur(ref_image, (5, 5), 0)
blurred_target = cv2.GaussianBlur(target_image, (5, 5), 0)

# 使用 Canny 算法检测边缘
edges_ref = cv2.Canny(blurred_ref, 100, 200)
edges_target = cv2.Canny(blurred_target, 100, 200)

# 创建一个三通道的图像用于显示结果
result_image = np.zeros((ref_image.shape[0], ref_image.shape[1], 3), dtype=np.uint8)

# 将参考图像的边缘用红色表示
result_image[edges_ref == 255] = [0, 0, 255]  # 红色

# 将目标图像的边缘用绿色表示
result_image[edges_target == 255] = [0, 255, 0]  # 绿色

# 将两幅图像的边缘重合部分用蓝色表示
both_edges = cv2.bitwise_and(edges_ref, edges_target)
result_image[both_edges == 255] = [255, 0, 0]  # 蓝色

# 显示结果
plt.figure(figsize=(12, 6))
plt.imshow(result_image)
plt.title("Edge Detection and Difference")
plt.axis('off')
plt.show()
